// Early Exception handling ft ..Rust
// By @5mukx

use std::{ffi::c_void, io::Read, ptr::null_mut};

use dinvk::{GetModuleHandle, GetProcAddress};

use crate::{
    hook::{hook_exception_dispatcher, set_hardware_breakpoint, unhook_exception_dispatcher},
    return_ssn::{NtAllocateVirtualMemoryFn, NtProtectVirtualMemoryFn},
};

mod debug;
mod hook;
mod return_ssn;
mod wow64_callback;

fn main() {
    let module = GetModuleHandle("NTDLL.DLL", None);

    if module.is_null() {
        println!("[ - ] Failed to get ntdll handle");
        return;
    }

    let module_base = module as *const u8;

    let wow64_prepare_for_exception = hook_exception_dispatcher(module_base);
    if wow64_prepare_for_exception.is_null() {
        println!("[ - ] Failed to hook exception dispatcher");
        return;
    }

    // get the address !
    let nt_alloc_addr = GetProcAddress(module, "NtAllocateVirtualMemory", None) as u64;
    let nt_protect_addr = GetProcAddress(module, "NtProtectVirtualMemory", None) as u64;

    if nt_alloc_addr == 0 || nt_protect_addr == 0 {
        println!("[ - ] Failed to get proc addresses");
        return;
    }

    unsafe {
        return_ssn::FUNCTION_ADDRESS_TABLE.nt_allocate_virtual_memory_address = nt_alloc_addr;
        return_ssn::FUNCTION_ADDRESS_TABLE.nt_protect_virtual_memory_address = nt_protect_addr;
    }


    let nt_allocate: NtAllocateVirtualMemoryFn = unsafe { std::mem::transmute(nt_alloc_addr) };
    let nt_protect: NtProtectVirtualMemoryFn = unsafe { std::mem::transmute(nt_protect_addr) };

    set_hardware_breakpoint(nt_alloc_addr);

    let mut base_address: *mut c_void = null_mut();
    let mut size: usize = 0x1000;

    let status = unsafe {
        nt_allocate(
            -1isize as *mut c_void,
            &mut base_address,
            0,
            &mut size,
            4096 | 8192, // mem commit | mem reserv
            4, // page read write
        )
    };

    if status != 0 {
        println!(
            "[ - ] NtAllocateVirtualMemory failed with status: 0x{:08X}",
            status
        );
        let _ = unhook_exception_dispatcher(wow64_prepare_for_exception);
        return;
    }

    println!("[ + ] Memory allocated at address: {:p}", base_address);

    set_hardware_breakpoint(nt_protect_addr);

    let mut old_protect: u32 = 0;
    let status = unsafe {
        nt_protect(
            -1isize as *mut c_void,
            &mut base_address,
            &mut size,
            32, // page exec reads
            &mut old_protect,
        )
    };

    if status != 0 {
        println!(
            "[ - ] NtProtectVirtualMemory failed with status: 0x{:08X}",
            status
        );
        let _ = unhook_exception_dispatcher(wow64_prepare_for_exception);
        return;
    }

    println!(
        "[ + ] Memory protection changed at address: {:p}",
        base_address
    );

    if !unhook_exception_dispatcher(wow64_prepare_for_exception) {
        println!("[ - ] Failed to unhook exception dispatcher");
        return;
    }

    // Pause to keep the console open, similar to getchar()
    let mut input = [0u8];
    let _ = std::io::stdin().read(&mut input);
}
